//https://leetcode.cn/problems/maximum-elegance-of-a-k-length-subsequence/?envType=daily-question&envId=2024-06-13

bool fun(vector<int>& v1, vector<int> v2){
    return v1[0]>v2[0];
}
class Solution {
public:
	long long findMaximumElegance(vector<vector<int>>& items, int k) {
		sort(items.begin(), items.end(), fun);
		long long total_profit = 0;     // 利润和
		long long distinct = 0;         // 类别和
		map<int, int> category_num;      // 类别
		vector<vector<int>> vv;          // 选择的数组
		for (int i = 0; i < k; i++) {
			vv.push_back(items[i]);
			total_profit += items[i][0];
			category_num[items[i][1]]++;
		}
		distinct = category_num.size();
        long long ret = total_profit + distinct * distinct;
		for (int i = k; i < items.size(); i++) {
			if (category_num[items[i][1]] == 0) {
					for (int ii = k - 1; ii >= 0; ii--) {
						if (category_num[vv[ii][1]] > 1) {

							long long newnum = total_profit - vv[ii][0] + items[i][0] + (distinct + 1) * (distinct + 1);
							ret = max(newnum, ret);
								total_profit = total_profit - vv[ii][0] + items[i][0];
								distinct++;
								category_num[vv[ii][1]]--;
								category_num[items[i][1]]++;
								swap(vv[ii], items[i]);
								break;
						}
					}
			}
            if(distinct == k){
                break;
            }
		}
		
		return  ret;
	}

};